package com.jerome.fs.mq;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Message Queue
 * @param <T> Type parameter for messages that are comparable
 * @author Jerome
 * @date 2023/9/7
 */
public class MessageQueue<T extends Comparable> {
    /**
     * Set of active topic names
     */
    private Set<String> activeTopics = new HashSet<>();
    /**
     * All topics
     */
    private Map<String, PriorityQueue<T>> topics;
    /**
     * Maximum capacity supported by the message queue
     */
    private int maxCapacity = Integer.MAX_VALUE;
    /**
     * Data counter, records the number of pending messages in the queue
     */
    private AtomicInteger count = new AtomicInteger(0);

    public MessageQueue() {
        this.topics = new HashMap<>();
    }

    public MessageQueue(int maxCapacity) {
        this.maxCapacity = maxCapacity;
        this.topics = new HashMap<>();
    }

    /**
     * Publish a message to a topic
     * @param topic The topic to which the message is published
     * @param content The content of the message
     * @throws Exception Throws an exception if the queue exceeds its maximum capacity
     */
    public synchronized void publish(String topic, T content) throws Exception {
        if (count.getAndIncrement() > maxCapacity) {
            throw new Exception("Too many requests");
        }
        topics.computeIfAbsent(topic, k -> new PriorityQueue<>()).offer(content);
        activeTopics.add(topic);
    }

    /**
     * Subscribe and retrieve a message from a topic
     * @param topic The topic to subscribe to
     * @return The retrieved message
     */
    public synchronized T subscribe(String topic) {
        T poll = topics.getOrDefault(topic, new PriorityQueue<>()).poll();
        if (poll == null) {
            activeTopics.remove(topic);
        }
        count.decrementAndGet();
        return poll;
    }

    /**
     * Get the active topics in the queue
     * @return Set of active topics
     */
    public Set<String> activeTopics() {
        return activeTopics;
    }

    /**
     * Get all topics in the queue
     * @return Set of all topics
     */
    public Set<String> topics() {
        return topics.keySet();
    }
}
